import torch
from transformers import AutoTokenizer, AutoModelForTokenClassification, AutoModelForSequenceClassification, pipeline
import re
import os

# --- 0. 配置和初始化 ---

# 假设这是从“20种药品.pdf”中提取的文本内容
pdf_content = """
了解居家常备药物：20 种常见类型的功能与用法概览

在日常生活中，面对突发的小病小痛，了解一些常见药物的基本知识能帮助我们更从容地应对。然而，药物是双刃剑，
知晓其功能和正确用法才能确保安全有效。本文将概述 20 种常见的药物类型，帮助您初步认识它们的主要作用和使用
方法。

首先，我们常备的药物中少不了解热镇痛药。这类药物又可细分为几类，例如非甾体抗炎药类，如布洛芬或阿司匹林，
它们主要用于缓解轻度到中度的疼痛，比如头痛、牙痛、肌肉痛和关节痛，同时具有退烧和抗炎的作用，通常口服使用，
需注意对胃肠道的潜在刺激和可能的过敏反应。另一类是对乙酰氨基酚类，如扑热息痛，这类药物的主要功能是退烧和
缓解轻度疼痛，对胃肠道的刺激相对较小，同样需要严格按照说明书的剂量口服，避免过量导致肝损伤。

当感冒引起咳嗽时，根据咳嗽的性质不同，会选择不同的药物。镇咳药，如右美沙芬，主要作用是抑制咳嗽反射，适用
于那些没有痰或者痰量很少的干咳和刺激性咳嗽，通过口服来缓解症状，但有痰的咳嗽通常不建议使用，以免影响痰液
排出。而对于痰多黏稠、咳痰困难的情况，则需要使用祛痰药，例如氨溴索、愈创甘油醚等，它们能够稀释痰液，帮助
痰液更容易咳出，通常也是口服使用，服用期间宜多饮水以助药效。

感冒或过敏常伴有鼻塞流涕，这时减充血剂就能派上用场。口服的减充血剂，如伪麻黄碱，通过收缩鼻腔血管来减轻鼻
塞和流涕症状，常用于缓解普通感冒或过敏引起的鼻塞，但需注意可能引起心率加快和血压升高，高血压或心脏病患者
需慎用。另一种是鼻腔喷雾类减充血剂，如赛洛唑啉，它们直接作用于鼻腔黏膜血管，能快速缓解鼻塞，通过鼻腔喷雾
使用，但这类药物不适合长期连续使用，否则可能导致药物性鼻炎。

过敏症状令人难受，这时就需要抗组胺药。第一代抗组胺药，如扑尔敏或苯海拉明，能有效缓解鼻痒、流涕、喷嚏以及
皮肤瘙痒等过敏症状，但通常伴有明显的镇静作用，口服使用，服药后应避免驾驶或操作精细机械。相比之下，第二代
抗组胺药，如氯雷他定或西替利嗪，也能缓解过敏症状，但镇静作用较弱或几乎没有，口服后对日常活动影响较小，适
合需要长时间控制过敏症状的人群。

消化不良、胃灼热和反酸是常见问题，可以选用抗酸药来快速缓解。铝碳酸镁或氢氧化铝等抗酸药通过中和胃酸起效，
能迅速缓解胃灼热、反酸和胃痛，通常以咀嚼片或混悬液形式口服，在胃部不适时服用，但效果比较短暂，不适合长期
大量使用。作用时间比抗酸药更长的是 H2 受体拮抗剂，如雷尼替丁或法莫替丁（请注意部分此类药物已转为处方药或
退市，但作为药物类型仍有认知度），它们通过减少胃酸分泌来缓解胃灼热和消化不良，口服使用，起效相对慢一些。
更强效抑制胃酸分泌的是质子泵抑制剂，部分种类如奥美拉唑已有非处方药，强效缓解胃酸相关症状，如频繁胃灼热，
通常需要连续服用几天效果最佳，不应作为日常助消化药物长期使用。

便秘是困扰不少人的问题，可以根据情况选择不同类型的轻泻药。容积性或膳食纤维类轻泻药，如欧车前，通过吸收水
分膨胀增加粪便体积，刺激肠道蠕动从而缓解便秘，口服使用，需配合大量水分，起效相对慢，通常需 1 到 3 天。渗透
性轻泻药，如聚乙二醇，能将水分保留在肠道内，软化粪便促进排便，通常溶解后口服，效果比较温和。而刺激性轻泻
药，如比沙可啶或番泻叶，直接刺激肠壁神经增加肠道蠕动，起效快，但可能引起腹痛和腹泻，口服使用，不宜长期或
频繁使用，以免产生依赖。

腹泻时若非感染引起，可以考虑使用止泻药，如洛哌丁胺，它能减慢肠道蠕动，减少排便次数，口服使用，但如果是细
菌性痢疾等感染性腹泻，则不应使用止泻药，以免延误病情。

旅行中容易遇到的问题是晕车、晕船、晕机，这时晕车药/晕船药就能发挥作用，如茶苯海明，它能抑制前庭功能，缓
解恶心、呕吐、头晕等症状，通常在出发前一定时间口服，这类药物常带有镇静作用。

家庭药箱中也少不了外用消毒剂或杀菌剂，如碘伏或酒精，它们用于清洁和消毒皮肤伤口、表皮擦伤等，防止感染，通
过外用在伤口周围皮肤，使用时需注意浓度和适用范围。皮肤瘙痒时，可以选用外用止痒药。皮质类固醇类，如氢化可
的松软膏，能缓解皮肤炎症和瘙痒，适用于湿疹、皮炎等，外用涂抹于患处，非处方药通常浓度较低，不宜长期大面积
使用。另一类外用止痒药是炉甘石洗剂等，具有收敛、保护和止痒作用，适用于痱子、轻度晒伤、蚊虫叮咬等，外用涂
抹于患处。

最后，对于因腹泻、呕吐或大量出汗导致的水分和电解质丢失，口服补液盐是重要的补充剂，能补充水分和电解质，预
防和纠正脱水，溶于水后口服，适用于轻中度脱水情况。

了解这些常见药物类型有助于更好地应对日常健康小问题，但务必记住，每种药物都有其特定的用法、用量、禁忌症和
注意事项。使用前务必仔细阅读药品说明书。如果症状持续不缓解、加重，或者您有基础疾病、正在服用其他药物、处
于怀孕或哺乳期，请务必及时咨询医生或药师。安全用药，健康生活。
"""

# 从HTML中提取的知识图谱Schema
# 允许的节点类型 (实体类型)
ALLOWED_NODE_TYPES = [
    "DrugType", "SpecificDrug", "Function", "ApplicableSymptom", "Usage", "Precaution"
]
# 允许的关系类型
ALLOWED_RELATION_TYPES = [
    "IS_A_TYPE_OF", "HAS_MAIN_FUNCTION", "APPLIES_TO", "HAS_USAGE",
    "HAS_PRECAUTION", "INCLUDES_EXAMPLE"
]

# 设置设备 (GPU优先)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device_name = 'cuda' if device.type == 'cuda' else 'cpu'
print(f"使用的设备: {device_name.upper()}")


# --- 1. 文本分句 ---
def split_into_sentences(text):
    sentences = re.split(r'[。？！]\s*[\n]?', text)
    return [s.strip() for s in sentences if s.strip()]


sentences = split_into_sentences(pdf_content)
print(f"文本分句完成，共 {len(sentences)} 句话。")

# --- 2. NER（命名实体识别）- 使用 iioSnail/bert-base-chinese-medical-ner (AutoModelForTokenClassification) ---
ner_model_name = "iioSnail/bert-base-chinese-medical-ner"

print(f"\n正在加载NER模型: {ner_model_name}...")
tokenizer_ner = AutoTokenizer.from_pretrained(ner_model_name)
model_ner = AutoModelForTokenClassification.from_pretrained(ner_model_name)
model_ner.to(device)  # 将模型移到设备上
model_ner.eval()  # 设置为评估模式

# 模型卡中给出的标签数字含义
# 1=B, 2=I, 3=E, 4=O
# 这里的labels是模型输出的ID到标签的映射
NER_ID_TO_TAG = {1: 'B', 2: 'I', 3: 'E', 4: 'O'}  # 假设模型内部的实际标签映射
# 并且模型输出就是这些数字
# 由于这个模型不提供实体类型，所以我们不能直接从模型输出中获取类型
# 而是识别出实体后，通过规则/词典进行类型补充

# --- 2.1 实体识别后处理及类型补充函数 (基于词典，与通用BERT模型时的逻辑类似) ---
# **重要：这些词典是识别实体和赋予类型的基础，请根据您的文本内容和Schema进行完善**
DRUGS_DICT = ["布洛芬", "阿司匹林", "扑热息痛", "右美沙芬", "氨溴索", "愈创甘油醚",
              "伪麻黄碱", "赛洛唑啉", "扑尔敏", "苯海拉明", "氯雷他定", "西替利嗪",
              "铝碳酸镁", "氢氧化铝", "雷尼替丁", "法莫替丁", "奥美拉唑",
              "欧车前", "聚乙二醇", "比沙可啶", "番泻叶", "洛哌丁胺", "茶苯海明",
              "碘伏", "酒精", "氢化可的松软膏", "炉甘石洗剂", "口服补液盐"]

SYMPTOMS_DISEASES_DICT = ["疼痛", "头痛", "牙痛", "肌肉痛", "关节痛", "过敏反应",
                          "咳嗽", "干咳", "刺激性咳嗽", "痰多黏稠", "咳痰困难",
                          "鼻塞", "流涕", "药物性鼻炎", "鼻痒", "喷嚏", "皮肤瘙痒",
                          "消化不良", "胃灼热", "反酸", "胃痛", "肝损伤", "便秘",
                          "腹泻", "细菌性痢疾", "感染性腹泻", "晕车", "晕船", "晕机",
                          "恶心", "呕吐", "头晕", "皮肤炎症", "湿疹", "皮炎", "痱子",
                          "晒伤", "蚊虫叮咬", "脱水"]

DRUG_TYPES_DICT = ["解热镇痛药", "非甾体抗炎药类", "对乙酰氨基酚类", "镇咳药", "祛痰药",
                   "减充血剂", "抗组胺药", "抗酸药", "H2 受体拮抗剂", "质子泵抑制剂",
                   "轻泻药", "容积性或膳食纤维类轻泻药", "渗透性轻泻药", "刺激性轻泻药",
                   "止泻药", "晕车药", "晕船药", "晕机药", "消毒剂", "杀菌剂",
                   "外用止痒药", "皮质类固醇类"]

# 这些关键词辅助判断实体的类型，但本身不直接作为实体
FUNCTION_TERM_MAP = {
    "缓解": "Function", "退烧": "Function", "抗炎": "Function", "抑制": "Function",
    "稀释": "Function", "帮助": "Function", "减轻": "Function", "收缩": "Function",
    "中和": "Function", "减少": "Function", "补充": "Function", "预防": "Function",
    "纠正": "Function", "收敛": "Function", "保护": "Function"
}
USAGE_TERM_MAP = {
    "口服": "Usage", "鼻腔喷雾": "Usage", "咀嚼片": "Usage", "混悬液": "Usage",
    "外用涂抹": "Usage", "溶于水后口服": "Usage", "外用在伤口周围": "Usage", "服用": "Usage",
    "使用": "Usage"
}
PRECAUTION_TERM_MAP = {
    "刺激": "Precaution", "过敏反应": "Precaution", "心率加快": "Precaution", "血压升高": "Precaution",
    "镇静作用": "Precaution", "肝损伤": "Precaution", "腹痛": "Precaution", "腹泻": "Precaution",
    "依赖": "Precaution", "禁忌症": "Precaution", "注意事项": "Precaution"
}


def assign_entity_type_by_dict(word, sentence_context=None):
    """
    通过词典规则为识别到的实体补充类型。
    """
    if word in DRUGS_DICT:
        return "SpecificDrug"
    elif word in SYMPTOMS_DISEASES_DICT:
        return "ApplicableSymptom"
    elif word in DRUG_TYPES_DICT:
        return "DrugType"
    elif word in FUNCTION_TERM_MAP:
        return FUNCTION_TERM_MAP[word]
    elif word in USAGE_TERM_MAP:
        return USAGE_TERM_MAP[word]
    elif word in PRECAUTION_TERM_MAP:
        return PRECAUTION_TERM_MAP[word]

    if word == "作用":
        return "Function"
    if word == "用法":
        return "Usage"
    if word == "注意事项":
        return "Precaution"

    return None


def convert_outputs_to_entities(sentence_tokens_ids, predicted_ids, sentence_text, tokenizer):
    """
    将模型输出的B/I/E/O ID序列转换为实体列表，并通过词典补充实体类型。
    Args:
        sentence_tokens_ids (list): 原始分词器分词后的tokens ID (包含特殊token CLS/SEP)。
        predicted_ids (list): 模型预测的每个token的B/I/E/O数字ID序列 (与 sentence_tokens_ids 对齐)。
        sentence_text (str): 原始句子文本。
        tokenizer: 用于获取token id和wordpiece信息的tokenizer。
    Returns:
        list: 识别出的实体列表，每个实体包含 text, type, start_char, end_char。
    """
    entities = []

    # 移除 [CLS] 和 [SEP] 对应的预测ID，以及填充ID
    # inputs['input_ids'][0] 是包含 [CLS] 和 [SEP] 的，predicted_ids_list 也对应
    # 移除特殊token对应的标签，只保留与句子实际token对应的部分

    # 获取有效长度 (去除CLS/SEP和padding)
    # 假设 attention_mask 已经将 padding 部分置为 0
    actual_len = sum(1 for x in predicted_ids if x != 0)  # 统计非O/非填充的标签
    if actual_len > 0 and predicted_ids[0] != 0:  # 如果CLS不是O，则实际长度是整个序列
        actual_len = len(predicted_ids)  # 此时predicted_ids已经去除了padding
    else:  # 否则，找到第一个非0的索引作为开始，最后一个非0的索引作为结束
        start_idx = 0
        for i, val in enumerate(predicted_ids):
            if val != 0:
                start_idx = i
                break
        end_idx = len(predicted_ids) - 1
        for i in range(len(predicted_ids) - 1, -1, -1):
            if predicted_ids[i] != 0:
                end_idx = i
                break
        actual_len = end_idx - start_idx + 1

    # 根据模型卡示例，outputs.logits.argmax(-1) * inputs['attention_mask']
    # 意味着 padding 部分的预测ID会是0。
    # 原始的 input_ids 包含了 [CLS] 和 [SEP]

    # 实际有效token和标签，去除CLS/SEP和padding
    # 这里的 indices 对应原始句子中的 token

    # 重新构建 token 和 tag 序列，跳过 CLS/SEP 和 padding (如果它们被标记为0)
    processed_tokens = []
    processed_tags_id = []
    # 从 input_ids_tensor[0] 和 predicted_ids_list 中构建
    for i in range(len(sentence_tokens_ids)):
        if sentence_tokens_ids[i] == tokenizer.cls_token_id or \
                sentence_tokens_ids[i] == tokenizer.sep_token_id or \
                sentence_tokens_ids[i] == tokenizer.pad_token_id or \
                predicted_ids[i] == 0:  # 跳过特殊token和padding，以及预测为O的token
            continue
        processed_tokens.append(tokenizer.convert_ids_to_tokens([sentence_tokens_ids[i]])[0])
        processed_tags_id.append(predicted_ids[i])

    current_entity_text = ""
    current_entity_tag_type = ""  # B, I, E, O
    current_entity_start_char = -1
    current_entity_end_char = -1

    current_offset_in_sentence = 0

    for i, (token, tag_id) in enumerate(zip(processed_tokens, processed_tags_id)):
        tag = NER_ID_TO_TAG[tag_id]

        original_token_segment = token
        if token.startswith('##'):
            original_token_segment = token[2:]

        token_start_in_sentence = sentence_text.find(original_token_segment, current_offset_in_sentence)

        if token_start_in_sentence == -1:
            current_offset_in_sentence += len(original_token_segment)
            continue
        token_end_in_sentence = token_start_in_sentence + len(original_token_segment)

        if tag == 'B':
            if current_entity_text:  # 结束上一个实体
                mapped_type = assign_entity_type_by_dict(current_entity_text, sentence_text)
                if mapped_type and mapped_type in ALLOWED_NODE_TYPES:
                    entities.append({
                        "text": current_entity_text,
                        "type": mapped_type,
                        "start_char": current_entity_start_char,
                        "end_char": current_entity_end_char,
                        "sentence_text": sentence_text
                    })
            current_entity_text = original_token_segment
            current_entity_tag_type = tag  # 记录B
            current_entity_start_char = token_start_in_sentence
            current_entity_end_char = token_end_in_sentence
        elif tag == 'I':
            if current_entity_text and (current_entity_tag_type == 'B' or current_entity_tag_type == 'I'):
                current_entity_text += original_token_segment
                current_entity_end_char = token_end_in_sentence
            else:  # I-标签但无B-开头或已结束，视为新的B-（或O）
                # 这种情况通常是模型输出不规则，尝试作为新实体
                if current_entity_text:  # 结束上一个实体
                    mapped_type = assign_entity_type_by_dict(current_entity_text, sentence_text)
                    if mapped_type and mapped_type in ALLOWED_NODE_TYPES:
                        entities.append({
                            "text": current_entity_text,
                            "type": mapped_type,
                            "start_char": current_entity_start_char,
                            "end_char": current_entity_end_char,
                            "sentence_text": sentence_text
                        })
                current_entity_text = original_token_segment
                current_entity_tag_type = tag  # 记录I
                current_entity_start_char = token_start_in_sentence
                current_entity_end_char = token_end_in_sentence
        elif tag == 'E':  # 实体结束
            if current_entity_text and (current_entity_tag_type == 'B' or current_entity_tag_type == 'I'):
                current_entity_text += original_token_segment
                current_entity_end_char = token_end_in_sentence
                # 实体结束，立即添加
                mapped_type = assign_entity_type_by_dict(current_entity_text, sentence_text)
                if mapped_type and mapped_type in ALLOWED_NODE_TYPES:
                    entities.append({
                        "text": current_entity_text,
                        "type": mapped_type,
                        "start_char": current_entity_start_char,
                        "end_char": current_entity_end_char,
                        "sentence_text": sentence_text
                    })
                # 重置
                current_entity_text = ""
                current_entity_tag_type = ""
            else:  # E标签但无B/I开头，视为O
                current_entity_text = ""
                current_entity_tag_type = ""
        else:  # 'O' 标签 (或未识别的标签)
            if current_entity_text:  # 结束上一个实体
                mapped_type = assign_entity_type_by_dict(current_entity_text, sentence_text)
                if mapped_type and mapped_type in ALLOWED_NODE_TYPES:
                    entities.append({
                        "text": current_entity_text,
                        "type": mapped_type,
                        "start_char": current_entity_start_char,
                        "end_char": current_entity_end_char,
                        "sentence_text": sentence_text
                    })
            current_entity_text = ""
            current_entity_tag_type = ""

        current_offset_in_sentence = token_end_in_sentence

        # 循环结束后，处理最后一个实体 (如果它是B或I结尾，没有E)
    if current_entity_text and (current_entity_tag_type == 'B' or current_entity_tag_type == 'I'):
        mapped_type = assign_entity_type_by_dict(current_entity_text, sentence_text)
        if mapped_type and mapped_type in ALLOWED_NODE_TYPES:
            entities.append({
                "text": current_entity_text,
                "type": mapped_type,
                "start_char": current_entity_start_char,
                "end_char": current_entity_end_char,
                "sentence_text": sentence_text
            })
    return entities


all_extracted_entities_per_sentence = []

print("\n正在进行NER识别...")
for i, sentence_text in enumerate(sentences):
    if not sentence_text:
        continue

    # 准备模型输入
    inputs = tokenizer_ner(
        [sentence_text],
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=tokenizer_ner.model_max_length,
        add_special_tokens=True  # 确保添加CLS/SEP
    )

    inputs = {k: v.to(device) for k, v in inputs.items()}

    try:
        with torch.no_grad():
            outputs = model_ner(**inputs)

        # 获取 logits，并转换为预测ID (argmax)
        # model_card 示例中是 outputs.logits.argmax(-1) * inputs['attention_mask']
        # 乘以 attention_mask 是为了将填充部分的预测ID设为0 (通常是O标签)
        predicted_ids_tensor = outputs.logits.argmax(dim=-1) * inputs['attention_mask']
        predicted_ids_list = predicted_ids_tensor.cpu().tolist()[0]  # 转换为list并取第一个句子

        # 将BERT分词器的input_ids传入，用于后处理
        current_sentence_entities = convert_outputs_to_entities(
            inputs['input_ids'][0].cpu().tolist(),  # 传入包含特殊token的input_ids
            predicted_ids_list,
            sentence_text,
            tokenizer_ner
        )

    except Exception as e:
        print(f"NER推理或后处理句子 '{sentence_text[:50]}...' 时发生错误: {e}")
        continue

    if current_sentence_entities:
        all_extracted_entities_per_sentence.append({
            "sentence_text": sentence_text,
            "entities": current_sentence_entities
        })

print(f"NER识别完成，共识别到 {sum(len(s['entities']) for s in all_extracted_entities_per_sentence)} 个实体。")

# --- 3. RE（关系抽取）- 模拟训练过的BERT分类器 (使用 bert-base-chinese) ---
re_base_model_name = "bert-base-chinese"

print(f"\n正在加载RE基础模型: {re_base_model_name}...")
tokenizer_re = AutoTokenizer.from_pretrained(re_base_model_name)

RE_LABELS = ["NO_REL"] + ALLOWED_RELATION_TYPES
RE_ID2LABEL = {i: label for i, label in enumerate(RE_LABELS)}
RE_LABEL2ID = {label: i for i, label in enumerate(RE_LABELS)}

model_re = AutoModelForSequenceClassification.from_pretrained(
    re_base_model_name,
    num_labels=len(RE_LABELS),
    id2label=RE_ID2LABEL,
    label2id=RE_LABEL2ID
)
model_re.eval()
if device.type == 'cuda':
    model_re.to(device)

special_tokens_for_re = ['<e1>', '</e1>', '<e2>', '</e2>']
tokenizer_re.add_tokens(special_tokens_for_re)
if len(tokenizer_re) > model_re.config.vocab_size:
    model_re.resize_token_embeddings(len(tokenizer_re))


# mark_entities_in_sentence 函数保持不变
def mark_entities_in_sentence(sentence, e1_span, e2_span, e1_text, e2_text):
    """
    在句子中用特殊标记 <e1></e1> 和 <e2></e2> 包裹两个实体。
    此版本尝试更精确地插入，并处理重叠情况
    """
    if sentence[e1_span[0]:e1_span[1]] != e1_text or sentence[e2_span[0]:e2_span[1]] != e2_text:
        e1_match = re.search(re.escape(e1_text), sentence)
        e2_match = re.search(re.escape(e2_text), sentence)
        if e1_match and e2_match:
            e1_span = e1_match.span()
            e2_span = e2_match.span()
        else:
            return None

    markers = [
        (e1_span[0], 0, '<e1>'), (e1_span[1], 1, '</e1>'),
        (e2_span[0], 0, '<e2>'), (e2_span[1], 1, '</e2>')
    ]
    markers.sort(key=lambda x: (x[0], -x[1], -len(x[2])))

    marked_sentence_chars = list(sentence)
    offset = 0
    for idx, marker_type, mark_text in markers:
        insert_pos = idx + offset
        try:
            marked_sentence_chars.insert(insert_pos, mark_text)
            offset += len(mark_text)
        except IndexError:
            print(
                f"警告: 尝试在 '{sentence}' 中插入标记 '{mark_text}' 时索引错误 at {insert_pos}/{len(marked_sentence_chars)}. Skipping.")
            return None

    return "".join(marked_sentence_chars)


# 定义关系类型过滤规则 (基于您的图谱Schema和允许的节点类型)
RELATION_TYPE_RULES = {
    ("DrugType", "SpecificDrug"): ["INCLUDES_EXAMPLE"],
    ("SpecificDrug", "Function"): ["HAS_MAIN_FUNCTION"],
    ("SpecificDrug", "ApplicableSymptom"): ["APPLIES_TO"],
    ("SpecificDrug", "Usage"): ["HAS_USAGE"],
    ("SpecificDrug", "Precaution"): ["HAS_PRECAUTION"],
    ("Function", "ApplicableSymptom"): ["APPLIES_TO"],
    # ... 您可以根据您的数据和Schema扩展这些规则
}

# --- 4. 关系抽取推理与三元组构建 (保持不变) ---
final_triples = []

print("\n正在进行关系抽取...")
for sentence_data in all_extracted_entities_per_sentence:
    sentence_text = sentence_data['sentence_text']
    entities_in_sentence = sentence_data['entities']

    for i in range(len(entities_in_sentence)):
        for j in range(len(entities_in_sentence)):
            if i == j:
                continue

            head_entity = entities_in_sentence[i]
            tail_entity = entities_in_sentence[j]

            candidate_relations = []
            if (head_entity['type'], tail_entity['type']) in RELATION_TYPE_RULES:
                candidate_relations.extend(RELATION_TYPE_RULES[(head_entity['type'], tail_entity['type'])])

            if not candidate_relations:
                continue

            marked_sentence = mark_entities_in_sentence(
                sentence_text,
                (head_entity['start_char'], head_entity['end_char']),
                (tail_entity['start_char'], tail_entity['end_char']),
                head_entity['text'], tail_entity['text']
            )

            if marked_sentence is None:
                continue

            inputs = tokenizer_re(marked_sentence, return_tensors="pt", truncation=True, max_length=512)

            if device.type == 'cuda':
                inputs = {k: v.to(device)}

            with torch.no_grad():
                outputs = model_re(**inputs)

            logits = outputs.logits
            predicted_relation_id = torch.argmax(logits, dim=-1).item()
            predicted_relation = RE_ID2LABEL[predicted_relation_id]

            if predicted_relation != "NO_REL" and predicted_relation in candidate_relations:
                triple = {
                    "head": head_entity['text'],
                    "head_type": head_entity['type'],
                    "relation": predicted_relation,
                    "tail": tail_entity['text'],
                    "tail_type": tail_entity['type'],
                    "source_sentence": sentence_text
                }
                if triple not in final_triples:
                    final_triples.append(triple)

print(f"\n三元组抽取完成，共抽取到 {len(final_triples)} 个三元组。")
print("\n--- 抽取到的三元组示例 (前20个) ---")
for i, triple in enumerate(final_triples[:20]):
    print(
        f"  {i + 1}. ({triple['head']}:{triple['head_type']}, {triple['relation']}, {triple['tail']}:{triple['tail_type']})")

